from machine import SPI, Pin, ADC
import utime
import st7789
import math

# ========== 引脚配置（根据实际连线修改） ==========
PIN_BUTTON_A = 0  # 物理按键 A 引脚
PIN_BUTTON_B = 14  # 物理按键 B 引脚
PIN_BUTTON_C = 9  # 物理按键 C 引脚
PIN_BUTTON_D = 10  # 物理按键 D 引脚

PIN_ADC_X = 12    # 摇杆 X (ADC)
PIN_ADC_Y = 13     # 摇杆 Y (ADC)
# ================================================

# ---------------- JoystickShield 类（来自你提供的库） ----------------
# 返回格式: [A, B, C, D, E, F]
# A-D: 1 未按下，0 按下
# E (索引4), F (索引5): 摇杆轴 -> -1 负方向, 0 中立, 1 正方向
class JoystickShield:
    def __init__(self, pinA, pinB, pinC, pinD, adc1, adc2):
        self.pinA = pinA
        self.pinB = pinB
        self.pinC = pinC
        self.pinD = pinD
        self.adc1 = adc1
        self.adc2 = adc2
        utime.sleep(0.1)

    def read(self):
        A = self.pinA.value()
        B = self.pinB.value()
        C = self.pinC.value()
        D = self.pinD.value()
        # 尝试设定均衡/采样模型（在部分平台上可用）
        try:
            self.adc1.equ(ADC.EQU_MODEL_8)
        except:
            pass
        try:
            E_raw = self.adc1.read()
        except:
            E_raw = 0
        try:
            self.adc2.equ(ADC.EQU_MODEL_8)
        except:
            pass
        try:
            F_raw = self.adc2.read()
        except:
            F_raw = 0
        # 根据经验阈值把ADC原始值映射为 -1/0/1
        E = -1 if E_raw < 850 else 1 if E_raw > 1050 else 0
        F = -1 if F_raw < 850 else 1 if F_raw > 1050 else 0
        # 如果两个轴同时都在同一极性（角落）则认为中立（避免误触）
        if (E == 1 and F == 1) or (E == -1 and F == -1):
            E, F = 0, 0
        data = [A, B, C, D, E, F]
        return data
# --------------------------------------------------------------------

# 创建物理按键 Pin 对象（先用 UP/DOWN/LEFT/RIGHT/BUTTON_A 名称表示物理引脚）
UP = Pin(PIN_BUTTON_A)
DOWN = Pin(PIN_BUTTON_B)
LEFT = Pin(PIN_BUTTON_C)
RIGHT = Pin(PIN_BUTTON_D)
BUTTON_A = Pin(PIN_BUTTON_A)  # 仅用于初始化设置，后面会重绑定为确认按钮对象

# 初始化按键
UP.init(Pin.IN)
DOWN.init(Pin.IN)
LEFT.init(Pin.IN)
RIGHT.init(Pin.IN)
BUTTON_A.init(Pin.IN)

UP.pull(Pin.PULL_DOWN)
DOWN.pull(Pin.PULL_DOWN)
LEFT.pull(Pin.PULL_DOWN)
RIGHT.pull(Pin.PULL_DOWN)
BUTTON_A.pull(Pin.PULL_DOWN)
# 上面是你指定的格式（注意：这里先用同名变量代表物理 Pin，后面会重绑定为逻辑按钮对象）

# 初始化 ADC（摇杆）
adcX = ADC(Pin(PIN_ADC_X))
adcY = ADC(Pin(PIN_ADC_Y))

# 创建 JoystickShield 实例，传入物理按键 Pin 对象
js = JoystickShield(UP, DOWN, LEFT, RIGHT, adcX, adcY)

# 虚拟按键包装：与游戏中原本的 UP.value()/DOWN.value()/... 调用兼容
class _DirButton:
    def __init__(self, js, axis_index, target_value):
        # axis_index: 4->X(E), 5->Y(F)
        self.js = js
        self.axis_index = axis_index
        self.target_value = target_value

    def value(self):
        data = self.js.read()
        try:
            axis = data[self.axis_index]
        except:
            return 0
        return 1 if axis == self.target_value else 0

class _ConfirmButton:
    def __init__(self, js):
        self.js = js

    def value(self):
        data = self.js.read()
        # A,B,C,D: 1 未按下，0 按下
        try:
            A, B, C, D = data[0], data[1], data[2], data[3]
        except:
            return 0
        pressed = (A == 0 or B == 0 or C == 0 or D == 0)
        # 游戏中使用 BUTTON_A.value()==1 表示确认按下
        return 1 if pressed else 0

# 创建四个方向和确认键对象，重绑定同名变量（覆盖之前的物理 Pin 变量）
# 轴索引：E (X) 在 read 返回中索引为 4，F (Y) 为 5
LEFT  = _DirButton(js, 4, -1)
RIGHT = _DirButton(js, 4, 1)
# 已调整：将 UP/DOWN 的 target_value 对调，使摇杆向上拨触发 UP
UP    = _DirButton(js, 5, 1)
DOWN  = _DirButton(js, 5, -1)
BUTTON_A = _ConfirmButton(js)  # 合并 ABCD 四键作为确认键

# =================== 以下为推箱子游戏代码 ===================

# 初始化SPI和屏幕（保持原来的引脚配置）
spi = SPI(0, baudrate=40000000, polarity=1, phase=0, bits=8, endia=0, sck=Pin(6), mosi=Pin(8))
display = st7789.ST7789(spi, 240, 240, reset=Pin(11, func=Pin.GPIO, dir=Pin.OUT), dc=Pin(7, func=Pin.GPIO, dir=Pin.OUT))
display.init()

# 颜色定义
BLACK = st7789.color565(0, 0, 0)
WHITE = st7789.color565(240, 240, 240)
RED = st7789.color565(255, 0, 0)
GREEN = st7789.color565(0, 255, 0)
BLUE = st7789.color565(0, 0, 255)
YELLOW = st7789.color565(255, 255, 0)
GRAY = st7789.color565(100, 100, 100)
PURPLE = st7789.color565(128, 0, 128)
ORANGE = st7789.color565(255, 165, 0)
CYAN = st7789.color565(0, 255, 255)
DARK_BLUE = st7789.color565(0, 0, 100)

# 游戏参数
GRID_SIZE = 20
GRID_WIDTH = 240 // GRID_SIZE
GRID_HEIGHT = 240 // GRID_SIZE

# 游戏状态
MENU = 0
GAME = 1
WIN = 2
LOSE = 3
SPLASH = 4
LEVEL_SELECT = 5

current_state = SPLASH
current_level = 0
current_difficulty = 0
difficulties = ["Easy", "Medium", "Hard"]
moves_count = 0
undo_stack = []
menu_selection = 0
level_selection = 0

# 关卡数据
levels = {
    "Easy": [
        {
            "map": [
                [1,1,1,1,1,1,1,1,1,1,1,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,1,1,1,1,0,0,0,0,1],
                [1,0,0,1,0,0,0,0,0,0,0,1],
                [1,0,0,1,0,0,0,0,0,0,0,1],
                [1,0,0,1,0,1,1,0,0,0,0,1],
                [1,0,0,1,0,0,1,0,0,0,0,1],
                [1,0,0,1,0,0,1,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,1,1,1,1,1,1,1,1,1,1,1]
            ],
            "player": (2,2),
            "boxes": [(4,4),(6,4),(7,7)],
            "targets": [(4,8),(8,4),(8,8)]
        },
        {
            "map": [
                [1,1,1,1,1,1,1,1,1,1,1,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,1,1,1,0,0,1],
                [1,0,0,1,1,0,0,0,1,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,1,0,0,0,0,0,1,0,0,1],
                [1,0,1,1,0,1,0,0,0,0,0,1],
                [1,0,0,0,0,1,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,1,1,1,1,1,1,1,1,1,1,1]
            ],
            "player": (2,2),
            "boxes": [(4,4),(5,5),(7,6)],
            "targets": [(8,4),(8,8),(4,8)]
        }
    ],
    "Medium": [
        {
            "map": [
                [1,1,1,1,1,1,1,1,1,1,1,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,1,1,1,1,0,0,1],
                [1,0,0,1,0,0,0,0,1,0,0,1],
                [1,0,0,1,0,0,0,0,0,0,0,1],
                [1,0,0,1,0,1,0,0,1,0,0,1],
                [1,0,0,1,0,1,0,0,0,0,0,1],
                [1,0,0,0,0,1,1,1,1,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,1,1,1,1,1,1,1,1,1,1,1]
            ],
            "player": (2,2),
            "boxes": [(4,4),(6,4),(7,6),(5,8)],
            "targets": [(8,4),(8,8),(4,8),(6,8)]
        }
    ],
    "Hard": [
        {
            "map": [
                [1,1,1,1,1,1,1,1,1,1,1,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,1,1,1,1,1,1,0,0,1],
                [1,0,0,1,0,0,0,0,1,0,0,1],
                [1,0,0,1,0,1,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,1,0,1,1,0,1,0,0,1],
                [1,0,0,1,0,0,0,0,1,0,0,1],
                [1,0,0,1,0,1,0,1,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,0,0,0,0,0,0,0,0,0,0,1],
                [1,1,1,1,1,1,1,1,1,1,1,1]
            ],
            "player": (2,2),
            "boxes": [(4,4),(6,4),(7,6),(5,7),(8,5)],
            "targets": [(8,4),(8,8),(4,8),(6,8),(5,5)]
        }
    ]
}

# 当前游戏状态
game_map = []
player = (0, 0)
boxes = []
targets = []

def save_state():
    return (player, boxes[:], moves_count)

def restore_state(state):
    global player, boxes, moves_count
    player, boxes, moves_count = state

def load_level(difficulty, level):
    global game_map, player, boxes, targets, moves_count, undo_stack
    level_data = levels[difficulties[difficulty]][level]
    game_map = [row[:] for row in level_data["map"]]
    player = level_data["player"]
    boxes = level_data["boxes"][:]
    targets = level_data["targets"][:]
    moves_count = 0
    undo_stack = []

def draw_game():
    display.fill(WHITE)
    for y in range(len(game_map)):
        for x in range(len(game_map[y])):
            if game_map[y][x] == 1:
                display.fill_rect(x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE - 1, GRID_SIZE - 1, BLACK)
    for (tx, ty) in targets:
        display.fill_rect(tx * GRID_SIZE, ty * GRID_SIZE, GRID_SIZE - 1, GRID_SIZE - 1, RED)
    for (bx, by) in boxes:
        box_color = PURPLE if (bx, by) in targets else BLUE
        display.fill_rect(bx * GRID_SIZE, by * GRID_SIZE, GRID_SIZE - 1, GRID_SIZE - 1, box_color)
    px, py = player
    display.fill_rect(px * GRID_SIZE, py * GRID_SIZE, GRID_SIZE - 1, GRID_SIZE - 1, GREEN)
    display.draw_string(10, 10, "Sokoban", color=YELLOW, size=2)
    display.draw_string(180, 10, "Moves: {}".format(moves_count), color=YELLOW, size=1)
    display.draw_string(10, 220, "Level: {}-{}".format(difficulties[current_difficulty], current_level + 1), color=YELLOW, size=1)

def check_win():
    for target in targets:
        if target not in boxes:
            return False
    return True

def show_win_screen():
    display.fill(BLACK)
    display.draw_string(70, 80, "YOU WIN!", color=GREEN, size=3)
    display.draw_string(60, 120, "Moves: {}".format(moves_count), color=YELLOW, size=2)
    display.draw_string(50, 160, "Press A to continue", color=WHITE, size=1)

def show_lose_screen():
    display.fill(BLACK)
    display.draw_string(70, 80, "GAME OVER", color=RED, size=3)
    display.draw_string(60, 120, "Boxes are stuck!", color=YELLOW, size=2)
    display.draw_string(50, 160, "Press A to continue", color=WHITE, size=1)

def show_hint(message):
    display.fill_rect(40, 100, 160, 40, BLACK)
    display.draw_string(50, 110, message, color=YELLOW, size=2)
    utime.sleep(1)
    draw_game()

def all_boxes_stuck():
    for bx, by in boxes:
        free_directions = 0
        for dx, dy in [(0,1),(0,-1),(1,0),(-1,0)]:
            new_x, new_y = bx + dx, by + dy
            if (0 <= new_x < GRID_WIDTH and 0 <= new_y < GRID_HEIGHT and
                    game_map[new_y][new_x] != 1 and (new_x, new_y) not in boxes):
                free_directions += 1
        if free_directions > 0:
            return False
    return True

def show_menu():
    display.fill(DARK_BLUE)
    display.draw_string(70, 30, "SOKOBAN", color=CYAN, size=4)
    # 只保留前两个选项：Play Game, Level Select
    options = ["Play Game", "Level Select"]
    for i, option in enumerate(options):
        color = ORANGE if i == menu_selection else WHITE
        display.draw_string(80, 100 + i * 30, option, color=color, size=2)
    display.draw_string(50, 210, "Use UP/DOWN and A to select", color=YELLOW, size=1)
    display.draw_string(10, 200, "by gyn", color=WHITE, size=1)  # 签名

def show_level_select():
    global level_selection
    display.fill(DARK_BLUE)
    display.draw_string(60, 20, "SELECT LEVEL", color=CYAN, size=3)

    # 显示难度列表（左侧）
    display.draw_string(30, 60, "Difficulty:", color=YELLOW, size=1)
    for i, diff in enumerate(difficulties):
        color = ORANGE if i == current_difficulty else WHITE
        display.draw_string(40, 80 + i * 28, diff, color=color, size=2)

    # 获取当前难度的关卡数并显示（右侧，起始 Y 偏移更低，避免与难度部分重叠）
    level_count = len(levels.get(difficulties[current_difficulty], []))
    if level_count == 0:
        display.draw_string(120, 120, "No levels", color=WHITE, size=2)
        display.draw_string(40, 210, "Use LEFT/RIGHT to change difficulty", color=YELLOW, size=1)
        display.draw_string(10, 200, "by gyn", color=WHITE, size=1)
        return

    # 保持 level_selection 在合法区间
    if level_selection >= level_count:
        level_selection = 0
    if level_selection < 0:
        level_selection = 0

    display.draw_string(120, 60, "Levels:", color=YELLOW, size=1)
    for i in range(level_count):
        color = ORANGE if i == level_selection else WHITE
        display.draw_string(120, 90 + i * 26, "Level {}".format(i + 1), color=color, size=2)

    display.draw_string(10, 210, "Use UP/DOWN to choose level", color=YELLOW, size=1)
    display.draw_string(10, 200, "by gyn", color=WHITE, size=1)  # 签名

def show_splash_screen():
    for y in range(240):
        color_val = int(128 + 127 * math.sin(y / 20))
        color = st7789.color565(0, 0, color_val)
        display.fill_rect(0, y, 240, 1, color)
    for i in range(10):
        display.draw_string(70, 100, "SOKOBAN", color=st7789.color565(0, 255, 255), size=4)
        utime.sleep(0.1)
        display.draw_string(70, 100, "SOKOBAN", color=st7789.color565(0, 0, 0), size=4)
        utime.sleep(0.1)
    display.draw_string(70, 100, "SOKOBAN", color=CYAN, size=4)
    # 开机动画处署名
    display.draw_string(10, 200, "by gyn", color=WHITE, size=1)
    for i in range(240):
        display.fill_rect(0, 200, i, 10, ORANGE)
        utime.sleep(0.01)
    utime.sleep(1)

def game_loop():
    global player, boxes, moves_count, undo_stack, current_state
    direction = (0, 0)
    last_input_time = utime.ticks_ms()
    draw_game()
    while current_state == GAME:
        current_time = utime.ticks_ms()
        # 使用摇杆控制方向，ABCD 任意键作为返回/确认
        if UP.value() == 1:
            direction = (0, -1)
        elif DOWN.value() == 1:
            direction = (0, 1)
        elif LEFT.value() == 1:
            direction = (-1, 0)
        elif RIGHT.value() == 1:
            direction = (1, 0)
        elif BUTTON_A.value() == 1:
            current_state = MENU
            utime.sleep_ms(300)
            return
        else:
            direction = (0, 0)
            utime.sleep_ms(50)
            continue
        # 防止按键过快
        if utime.ticks_diff(current_time, last_input_time) < 150:
            utime.sleep_ms(50)
            continue
        last_input_time = current_time
        current_state_save = save_state()
        px, py = player
        dx, dy = direction
        new_px, new_py = px + dx, py + dy
        if 0 <= new_px < GRID_WIDTH and 0 <= new_py < GRID_HEIGHT:
            if game_map[new_py][new_px] == 1:
                continue
            if (new_px, new_py) in boxes:
                new_bx, new_by = new_px + dx, new_py + dy
                if (0 <= new_bx < GRID_WIDTH and 0 <= new_by < GRID_HEIGHT and
                        game_map[new_by][new_bx] != 1 and (new_bx, new_by) not in boxes):
                    boxes.remove((new_px, new_py))
                    boxes.append((new_bx, new_by))
                    undo_stack.append(current_state_save)
                else:
                    continue
            player = (new_px, new_py)
            moves_count += 1
            draw_game()
            if check_win():
                show_win_screen()
                current_state = WIN
                utime.sleep_ms(500)
                return
            if all_boxes_stuck():
                show_lose_screen()
                current_state = LOSE
                utime.sleep_ms(500)
                return
        utime.sleep_ms(50)

def main():
    global current_state, menu_selection, level_selection, current_difficulty, current_level
    show_splash_screen()
    utime.sleep(2)
    current_state = MENU
    while True:
        if current_state == MENU:
            show_menu()
            while current_state == MENU:
                # 只保留两个选项，所以使用 2 做模
                if UP.value() == 1:
                    menu_selection = (menu_selection - 1) % 2
                    show_menu()
                    utime.sleep_ms(200)
                elif DOWN.value() == 1:
                    menu_selection = (menu_selection + 1) % 2
                    show_menu()
                    utime.sleep_ms(200)
                elif BUTTON_A.value() == 1:
                    if menu_selection == 0:
                        load_level(current_difficulty, current_level)
                        current_state = GAME
                    elif menu_selection == 1:
                        # 进入关卡选择
                        level_selection = 0
                        current_state = LEVEL_SELECT
                    utime.sleep_ms(300)
                utime.sleep_ms(50)
        elif current_state == LEVEL_SELECT:
            show_level_select()
            while current_state == LEVEL_SELECT:
                # 读取当前难度的关卡数量
                level_count = len(levels.get(difficulties[current_difficulty], []))
                # 上/下 控制关卡的选择
                if UP.value() == 1 and level_count > 0:
                    level_selection = (level_selection - 1) % level_count
                    show_level_select()
                    utime.sleep_ms(200)
                elif DOWN.value() == 1 and level_count > 0:
                    level_selection = (level_selection + 1) % level_count
                    show_level_select()
                    utime.sleep_ms(200)
                # 左/右 切换难度（左为上一个，右为下一个）
                elif LEFT.value() == 1:
                    current_difficulty = (current_difficulty - 1) % len(difficulties)
                    # 更新 level_count 与 level_selection
                    level_selection = 0
                    show_level_select()
                    utime.sleep_ms(200)
                elif RIGHT.value() == 1:
                    current_difficulty = (current_difficulty + 1) % len(difficulties)
                    level_selection = 0
                    show_level_select()
                    utime.sleep_ms(200)
                elif BUTTON_A.value() == 1:
                    # 防抖：等待按键释放后再确认，避免误触立即进入
                    while BUTTON_A.value() == 1:
                        utime.sleep_ms(30)
                    utime.sleep_ms(80)
                    level_count = len(levels.get(difficulties[current_difficulty], []))
                    if level_count == 0:
                        # 没有关卡则提示并返回菜单
                        show_hint("No levels")
                        current_state = MENU
                    else:
                        current_level = level_selection
                        load_level(current_difficulty, current_level)
                        current_state = GAME
                    utime.sleep_ms(300)
                utime.sleep_ms(50)
        elif current_state == GAME:
            game_loop()
        elif current_state == WIN:
            # 等待按键动作，采用先等待释放再等待按下的方式以避免误触
            while BUTTON_A.value() == 1:
                utime.sleep_ms(50)
            while BUTTON_A.value() == 0:
                utime.sleep_ms(50)
            # 按下后自动尝试进入下一关，如果没有下一关则返回菜单
            level_count = len(levels[difficulties[current_difficulty]])
            if current_level + 1 < level_count:
                current_level += 1
                load_level(current_difficulty, current_level)
                current_state = GAME
                utime.sleep_ms(300)
            else:
                # 没有更多关卡，返回菜单
                current_state = MENU
                utime.sleep_ms(300)
        elif current_state == LOSE:
            while BUTTON_A.value() == 0:
                utime.sleep_ms(50)
            current_state = MENU
            utime.sleep_ms(300)
        elif current_state == SPLASH:
            show_splash_screen()
            utime.sleep(2)
            current_state = MENU
        utime.sleep_ms(50)

if __name__ == "__main__":
    main()
